# -*-coding:utf-8 -*-

"""
# File       : filestest.py
# Time       ：2021/11/12 15:33
# Author     ：dcheng.z
# Description：
"""
import pandas as pd
import os,cv2
from sklearn.model_selection import train_test_split
import torchvision.transforms as transforms
from load_data import GAMMA_sub1_dataset
from torch.utils.data import DataLoader

dataset_root = '../GAMMA_training data/training_data/multi-modality_images'

label_file='../GAMMA_training data/training_data/glaucoma_grading_training_GT.xlsx'

class Test:
    """
    getitem() output:
        fundus_img: RGB uint8 image with shape (3, image_size, image_size)
        oct_img:    Uint8 image with shape (256, oct_img_size[0], oct_img_size[1])
    """

    def __init__(self,
                 dataset_root,
                 label_file='',
                 filelists=None,
                 num_classes=3,
                 mode = 'train'):

        self.dataset_root = dataset_root

        self.num_classes = num_classes
        label = {row['data']: row[1:].values for _, row in pd.read_excel(label_file).iterrows()}

        self.file_list = [[f, label[int(f)]] for f in os.listdir(dataset_root)]

        train_filelists, val_filelists = train_test_split(self.file_list, test_size=0.2, random_state=42)

        if filelists is not None:   # filelists=None
            self.file_list = [item for item in self.file_list if item[0] in filelists]

    def __getitem__(self, idx):
        real_index, label = self.file_list[idx]
        print("idx",idx)
        fundus_img_path = os.path.join(self.dataset_root, real_index, real_index + ".jpg")

        oct_series_list = sorted(os.listdir(os.path.join(self.dataset_root, real_index, real_index)),
                            key=lambda x: int(x.split("_")[0]))

        oct_series_0 = cv2.resize(cv2.imread(os.path.join(self.dataset_root, real_index, real_index, oct_series_list[0]),
                            cv2.IMREAD_GRAYSCALE),(256,256))
        print("idx",oct_series_list[idx])


image_size = 512
img_train_transforms = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomResizedCrop(image_size, scale=(0.8, 1)),
    transforms.Resize((image_size, image_size)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(contrast=0.4,brightness=0.4),

    # transforms.ColorJitter(contrast=0.4,brightness=0.4),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
])

oct_train_transforms = transforms.Compose([
#        transforms.ToPILImage(),
#        transforms.Resize((image_size, image_size)),
#        transforms.RandomHorizontalFlip(),
#        transforms.RandomRotation(15),
#        transforms.ColorJitter(contrast=0.2,brightness=0.2),
    transforms.ToTensor(),
#        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
])
trainset_root='../GAMMA_training data/training_data/multi-modality_images'

train_dataset = GAMMA_sub1_dataset(dataset_root=trainset_root, # 训练数据和val数据在文件夹进行分配
                        img_transforms=img_train_transforms,
                        oct_transforms=oct_train_transforms,
                        label_file='../GAMMA_training data/training_data/glaucoma_grading_training_GT.xlsx',
                        mode='train')
# print(fundus_img.shape,oct_img.shape,label.shape)

print(train_dataset)
batchsize = 8 # 批大小,
train_loader = DataLoader(train_dataset, batch_size=batchsize,pin_memory=True, num_workers=0,shuffle=False)
print(train_loader)
